from pdboi.pdbdata.mmcifapiconnector import MMcifAPIconnector, MMcifCategoryNotFoundException
from atomsitecache import AtomsiteCache, OrderedSet
from exceptions import AtomNotModelledException, PDBentityNotFoundException

from collections import OrderedDict

import string, sys, copy


class ModelledGroup :
	'''this class represents the set of atoms identified in PDB file by a model id, chain id, residue number, insertion code.
		Depending on altcode (ie split residue or microhet), there can be more than 1 subgroup. subgroup can be mapped to a ChemicalGroup.
	'''
	def __init__(s, dataconnector, modelid, chainid, resnum, icode) :
		s.dc, s.modelid, s.chainid, s.resnum, s.icode = dataconnector, modelid, chainid, resnum, icode
		s.entry_id = s.dc.getCategory("entry","id","first")
	def getId(s) : return (s.entry_id, s.modelid, s.chainid, s.resnum, s.icode)
	def getEntryId(s) : return s.entry_id
	def getModelId(s) : return s.modelid
	def getChainId(s) : return s.chainid
	def getResnum(s) : return s.resnum
	def getInsertcode(s) : return s.icode
	def getIcodedResnum(s) :
		ret = s.resnum
		if s.icode in [".","?"] : ret += " "
		else : ret += s.icode
		return ret
	def getSubgroups(s) :
		'''check if there is any splitting of group's atoms on account of microhet on altcode, make subgroups and return a list'''
		if not hasattr(s,'subgroups') : s.makeSubgroups()
		return s.subgroups.values()
	def getAltcodesForGroupname(s, name) :
		if not hasattr(s,'subgroups') : s.makeSubgroups()
		return [altcode for resname,altcode in s.subgroups.keys() if resname==name]
	def getSubgroup(s, resname, altcode) :
		'''return a particular subgroup or raise exception. altcodes ".","?"," " are considered identical!'''
		if not hasattr(s,'subgroups') : s.makeSubgroups()
		if altcode in ['.','?',' '] :
			for ac in ['.','?',' '] :
				try : return s.subgroups[(resname,ac)]
				except KeyError : pass
		try : return s.subgroups[(resname,altcode)]
		except KeyError : raise PDBentityNotFoundException("Subgroup with given resname,altcode not found: " + str(resname) + " and " + str(altcode) + " Have these instead : " + str(s.subgroups.keys()) )
	def makeSubgroups(s) :
		s.subgroups = OrderedDict()
		mcra = AtomsiteCache.getCache(s.dc).getHierarchy()
		resdict = mcra[s.modelid][s.chainid][(s.resnum,s.icode)]
		resname_altcodes = OrderedSet()
		for resname, atoms in resdict.items() :
			for atominfo in atoms :
				resname_altcodes.add( (resname,atominfo[1]) )
		todelete = []
		for resname,altcode in resname_altcodes :
			if altcode not in [" ","?","."] :
				for achar in ["."," ","?"] :
					if (resname,achar) in resname_altcodes : resname_altcodes.remove( (resname,achar) )
		for resname,altcode in resname_altcodes :
			s.subgroups[(resname,altcode)] = ModelledSubgroup(s.dc, s.modelid, s.chainid, s.resnum, s.icode, resname, altcode)

class ModelledSubgroup :
	'''This represents a modelled group identified in PDB file by a model id, chain id, residue number, insertion code, resname, altcode.
	Atoms identified with altcodes '.', '?', ' ' are considered part of this subgroup if everything else matches.'''
	def __init__(s, dataconnector, modelid, chainid, resnum, icode, resname, altcode) :
		s.dc, s.modelid, s.chainid, s.resnum, s.icode, s.resname, s.altcode = dataconnector, modelid, chainid, resnum, icode, resname, altcode
		s.entry_id = s.dc.getCategory("entry","id","first")
	def getId(s) : return (s.entry_id, s.modelid, s.chainid, s.resnum, s.icode, s.resname, s.altcode)
	def getEntryId(s) : return s.entry_id
	def getModelId(s) : return s.modelid
	def getChainId(s) : return s.chainid
	def getResnum(s) : return s.resnum
	def getInsertcode(s) : return s.icode
	def getAltcode(s) : return s.altcode
	def getResname(s) : return s.resname
	def getNumExpectedModelledAtoms(s) : raise NotImplementedError()
	def prepareCoordiates(s) :
		s.coords = AtomsiteCache.getCache(s.dc).getSubgroupCoords(s.getModelId(),s.getChainId(),s.getResnum(),s.getInsertcode(),s.getResname(),s.getAltcode())
	def getAtomCoordinate(s, atomid) :
		if not hasattr(s, "coords") or s.coords == None : s.prepareCoordiates()
		try : return s.coords[atomid]
		except KeyError : raise AtomNotModelledException(str(s.getId()) + " " + atomid)
	def getNumModelledAtoms(s) :
		if not hasattr(s, "coords") or s.coords == None : s.prepareCoordiates()
		return len(s.coords)

class ChemicalGroup :
	'''this class represents reference checmical group as described in chemical component dictionary.
		it also caches upto 500 instances of itself so that they are not created repeatedly.
	'''
	@staticmethod
	def getChemicalGroup(compname,chemcompDictpath="/ebi/msd/work2/msdsd/data/ligand-dict/ligand-dict-v3") :
		'''factory method, use instead of constructor to avoid unnecessary instances'''
		maxCachesize = 500
		if not hasattr(ChemicalGroup, "cache") : ChemicalGroup.cache = {}
		if not ChemicalGroup.cache.has_key(compname) and len(ChemicalGroup.cache) < maxCachesize :
			ChemicalGroup.cache[compname] = ChemicalGroup(compname,chemcompDictpath)
		return ChemicalGroup.cache[compname]
	def __init__(s, compname, chemcompDictpath) :
		dictpath = "%s/%s/%s/%s.cif" % (chemcompDictpath, compname[0], compname, compname)
		s.dc, s.compname = MMcifAPIconnector(dictpath), compname
	def getName(s) : return s.compname
	def prepareBonds(s) :
		s.bonds = OrderedDict()
		try :
			for bonddata in s.dc.getCategory("chem_comp_bond",None,"iter") :
				s.bonds[ (bonddata['atom_id_1'],bonddata['atom_id_2']) ] = ChemicalGroup.Bond( bonddata )
		except MMcifCategoryNotFoundException : pass # possible to have no bonds
	def prepareAtoms(s) :
		s.atoms = OrderedDict()
		for atomdata in s.dc.getCategory("chem_comp_atom",None,"iter") :
			s.atoms[ atomdata['atom_id'] ] = ChemicalGroup.Atom( atomdata )
	def getBonds(s) :
		if not hasattr(s, "bonds") : s.prepareBonds()
		return s.bonds.values()
	def numBonds(s) :
		if not hasattr(s, "bonds") : s.prepareBonds()
		return len(s.bonds)
	def getAtoms(s) :
		if not hasattr(s, "atoms") : s.prepareAtoms()
		return s.atoms.values()
	def numAtoms(s) :
		if not hasattr(s, "atoms") : s.prepareAtoms()
		return len(s.atoms)
	def getElement(s, atomname) :
		if not hasattr(s, "atoms") : s.prepareAtoms()
		return s.atoms[atomname].getElement()
	class Atom :
		def __init__(s, atomdata) : s.cifdata = atomdata
		def getElement(s) : return s.cifdata['type_symbol']
		def getAtomId(s) : return s.cifdata['atom_id']
		def isAromatic(s) : return s.cifdata['pdbx_aromatic_flag'] == "Y"
		def getStereo(s) : return s.cifdata['pdbx_stereo_config']
		def getFormalCharge(s) : return string.atoi(s.cifdata['charge'])
		def getIdealCoords(s) : return tuple([string.atof(s.cifdata['pdbx_model_Cartn_%s_ideal' % xyz]) for xyz in ['x','y','z']])
	class Bond :
		def __init__(s,bonddata) : s.cifdata = bonddata
		def getAtomId1(s) : return s.cifdata['atom_id_1']
		def getAtomId2(s) : return s.cifdata['atom_id_2']
		def getBondOrder(s) : return s.cifdata['value_order']
		def isAromatic(s) : return s.cifdata['pdbx_aromatic_flag'] == "Y"
		def getStereo(s) : return s.cifdata['pdbx_stereo_config']
	def writeMol2file(s,subgroup,mol2fn) :
		'''write a mol2 file to given path using coordinates from given subgroup using openbabel.
		unmodelled atoms are not written to mol2 file.'''
		import pybel, openbabel
		mol = s.makeOpenbabelMolecule(subgroup)
		obConversion = openbabel.OBConversion()
		obConversion.SetInAndOutFormats("mol2", "mol2")
		obConversion.WriteFile(mol, mol2fn)
	def makeOpenbabelMolecule(s, subgr=None, useIdealCrdsForHydrogens=None) :
		'''create openbabel OBMol. Use coordinates from subgr if given. Copy ideal coordinates of chemcomp dictionary for hydrogens if useIdealCrdsForHydrogens is set to 1.'''
		if not hasattr(s,'atoms') : s.prepareAtoms()
		if not hasattr(s,'bonds') : s.prepareBonds()
		import pybel, openbabel
		mol = openbabel.OBMol() ; oet = openbabel.OBElementTable()
		mol.BeginModify()
		mol.NewResidue() # need to create residue and add atoms to it also in order to keep name atoms in mol2
		unmodelledAtoms = [] ; atomid2obaid = {}
		for at in s.atoms.values() :
			if subgr != None :
				try : acrd = subgr.getAtomCoordinate( at.getAtomId() )
				except AtomNotModelledException, anm :
					if useIdealCrdsForHydrogens and at.getElement() in ['H','D'] : acrd = at.getIdealCoords()
					else : unmodelledAtoms.append(at.getAtomId()) ; continue
			a = mol.NewAtom() ; a.SetAtomicNum( oet.GetAtomicNum(at.getElement()) )
			if at.isAromatic() : a.SetAromatic()
			a.SetFormalCharge(at.getFormalCharge())
			if at.getStereo() == 'R' : a.SetPositiveStereo()
			elif at.getStereo() == 'S' : a.SetNegativeStereo()
			if subgr != None : a.SetVector( acrd[0], acrd[1], acrd[2] )
			mol.GetResidue(0).AddAtom(a) ; mol.GetResidue(0).SetAtomID(a,at.getAtomId())
			atomid2obaid[at.getAtomId()] = len(atomid2obaid) + 1
		bondorder2num = {'sing':1, 'doub':2, 'trip':3}
		for bn in s.bonds.values() :
			bt = bondorder2num[bn.getBondOrder().lower()]
			if bn.isAromatic() : bt = 5
			if bn.getAtomId2() in unmodelledAtoms or bn.getAtomId1() in unmodelledAtoms : continue
			obaid1, obaid2 = atomid2obaid[bn.getAtomId1()], atomid2obaid[bn.getAtomId2()]
			mol.AddBond( obaid1, obaid2, bt )
			if   bn.getStereo() == 'E' : mol.GetBond(obaid1,obaid2).SetUp()
			elif bn.getStereo() == 'Z' : mol.GetBond(obaid1,obaid2).SetDown()
		mol.EndModify()
		return mol
